import glob
import os.path
import pickle

import matplotlib.pyplot as plt
import numpy as np
import scipy.io.wavfile as wav
import cv2
import copy
import scipy.signal as signal

from crop_wav.crop_utils import *



class SingleWavShower:
    def __init__(self, filePath, params):
        self.filePath = filePath
        self.params = params
        self.sonicFreq = params.sonicFreq
        self.lipOffset = params.lipOffset
        self.nfft = params.nfft
        self.frameTime = params.frameTime
        self.aheadTime = params.aheadTime
        self.isGetRange = params.isGetRange

    def getRangeSTFTData(self, data, borderFreq):
        start = borderFreq[0]
        end = borderFreq[1]
        limits = getLipLimitIndex(start, end, self.fs, self.nfft)
        f, t, Zxx = data
        f = f[limits[0]:limits[1]]
        Zxx = Zxx[limits[0]:limits[1], :]
        data[0] = f
        data[2] = Zxx
        return data

    def getBorderFreq(self, sonicFreq, lipOffset):
        start = sonicFreq - lipOffset
        end = sonicFreq + lipOffset
        return [start, end]

    def getSTFTData(self, fs, originSignal):
        nperseg = int(fs * self.frameTime)
        noverlap = int(fs * (self.frameTime - self.aheadTime))
        f, t, Zxx = signal.stft(originSignal, fs, nperseg=nperseg, noverlap=noverlap, nfft=self.nfft)
        Zxx = np.abs(Zxx)
        return [f, t, Zxx]

    def getFinalCropLen(self):
        self.fs, originSignal = wav.read(self.filePath)
        data = self.getSTFTData(self.fs, originSignal)
        borderFreq = self.getBorderFreq(self.sonicFreq, self.lipOffset)
        data = self.getRangeSTFTData(data, borderFreq)
        f, t, Zxx = data
        Zxx = subAdjacent(Zxx)
        t = t[:Zxx.shape[1]]

        print(self.filePath)

        padNum = 20
        sumArr = np.sum(Zxx, axis=0)
        sumArr = sumArr ** 2
        sumArr = sumArr[padNum:len(sumArr) - padNum]

        avgArr = AvgZeroPad(sumArr, 20)
        avgArr = avgArr / np.max(avgArr)
        thresh = 0.1

        activation = getActivation(avgArr, thresh)
        clipIndex = getClipIndex(activation)
        concatIndex = getConcatIndex(clipIndex, 20)
        resIndex = getMaxLenIndex(concatIndex)

        return (resIndex[1] - resIndex[0])


    def getFinalCrop(self, lengthDict):
        self.fs, originSignal = wav.read(self.filePath)
        data = self.getSTFTData(self.fs, originSignal)
        borderFreq = self.getBorderFreq(self.sonicFreq, self.lipOffset)
        data = self.getRangeSTFTData(data, borderFreq)
        f, t, Zxx = data
        Zxx = subAdjacent(Zxx)
        t = t[:Zxx.shape[1]]

        padNum = 20
        sumArr = np.sum(Zxx, axis=0)
        sumArr = sumArr ** 2
        sumArr = sumArr[padNum:len(sumArr) - padNum]

        avgArr = AvgZeroPad(sumArr, 20)
        avgArr = avgArr / np.max(avgArr)
        thresh = 0.1

        activation = getActivation(avgArr, thresh)
        clipIndex = getClipIndex(activation)
        concatIndex = getConcatIndex(clipIndex, 20)
        resIndex = getMaxLenIndex(concatIndex)

        meanStd = getMeanAndStdOfWordLen(lengthDict)

        word = getWordFromPath(self.filePath)
        # cropLen = getCropLenOfOneWord(meanStd, word)
        cropLen = getCropLenOfAllWord(meanStd)
        finalIndex = getFinalIndexRange(resIndex, cropLen)
        finalIndex[1] = min(finalIndex[1], len(avgArr)-1-padNum)
        return [finalIndex[0]+padNum, finalIndex[1]+padNum]

    def getSingleCrop(self):
        self.fs, originSignal = wav.read(self.filePath)
        data = self.getSTFTData(self.fs, originSignal)
        borderFreq = self.getBorderFreq(self.sonicFreq, self.lipOffset)
        data = self.getRangeSTFTData(data, borderFreq)
        f, t, Zxx = data
        Zxx = subAdjacent(Zxx)
        t = t[:Zxx.shape[1]]

        print(self.filePath)

        padNum = 20
        sumArr = np.sum(Zxx, axis=0)
        sumArr = sumArr ** 2
        sumArr = sumArr[padNum:len(sumArr) - padNum]

        avgArr = AvgZeroPad(sumArr, 20)
        avgArr = avgArr / np.max(avgArr)
        thresh = 0.1

        activation = getActivation(avgArr, thresh)
        clipIndex = getClipIndex(activation)
        concatIndex = getConcatIndex(clipIndex, 20)
        resIndex = getMaxLenIndex(concatIndex)

        # plt_data = np.zeros(Zxx.shape[1])
        # plt_data[resIndex[0]:resIndex[1]] = 1
        # plt.plot(avgArr)
        # plt.plot(plt_data)
        # plt.show()
        resIndex[1] = min(resIndex[1], len(avgArr) - 1 - padNum)
        return [resIndex[0]+padNum, resIndex[1]+padNum]

class MultiWavShower:
    def __init__(self, dirPath, params):
        self.dirPath = dirPath
        self.params = params
        self.allPaths = self.getAllPath(dirPath)
        self.lenDictPath = params.lenDictPath
        self.cropMsgPath = params.cropMsgPath

    def getAllPath(self, dirPath):
        allPaths = []
        pattern = os.path.join(dirPath, '*.wav')
        for path in glob.glob(pattern):
            allPaths.append(path)
        allPaths.sort()
        return allPaths

    def getMaxLenOfEachWord(self):
        lengthDict = {}
        for path in self.allPaths:
            word = getWordFromPath(path)
            lengthDict[word] = []
        for path in self.allPaths:
            obj = SingleWavShower(path, self.params)
            curLen = obj.getFinalCropLen()
            word = getWordFromPath(path)
            lengthDict[word].append(curLen)
        print(lengthDict)
        saveLenDictPath = os.path.join(self.dirPath, self.lenDictPath)
        with open(saveLenDictPath, 'wb') as f:
            pickle.dump(lengthDict, f)

    def saveTimeIndexInFile(self):
        cropLenPath = os.path.join(self.dirPath, self.lenDictPath)
        # if os.path.exists(cropLenPath):
        #     with open(cropLenPath, 'rb') as f:
        #         lengthDict = pickle.load(f)
        # else:
        #     raise EOFError

        timeMap = {}
        for path in self.allPaths:
            obj = SingleWavShower(path, self.params)
            # cropIndex = obj.getFinalCrop(lengthDict)
            cropIndex = obj.getSingleCrop()
            finalPath = path.replace(os.sep, '/')
            finalPath = finalPath.split('/')[-1]
            timeMap[finalPath] = cropIndex
        saveTxt = os.path.splitext(self.cropMsgPath)[0] + '.txt'
        saveDp = self.cropMsgPath

        saveTxt = os.path.join(self.dirPath, saveTxt)
        saveDp = os.path.join(self.dirPath, saveDp)

        with open(saveTxt, 'w') as f:
            for key in timeMap.keys():
                msg = f'{key} {timeMap[key][0]} {timeMap[key][1]}\n'
                f.write(msg)

        with open(saveDp, 'wb') as f:
            pickle.dump(timeMap, f)

    def getFinalCropMsg(self):
        # self.getMaxLenOfEachWord()
        self.saveTimeIndexInFile()



